import logging
import numpy as np

from tensorflow.keras.models import Model, load_model

from util.util_log import setup_logging
from util.util_ris_pattern import point_2_phase, phase_2_pattern, phase_2_pattern_xyz, eps, \
    phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft

from dissertation.util_phase_pattern import traces_2_phaseRads, phases_rad_2_bit, phaseBit_2_pattern_point
from dissertation.util_plot import save_images_2x2, plot_images_2x2, save_images_3x2, plot_images_3x2

from multi_beam.multi_beam_PS_complex import psm_complex_beam_n


import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file=base_path + "/trace.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[beam-trace-autocoder]")

# ============================================= 读取自编码器相关 =======================================
def load_models(path_encoder, path_decoder):
    encoder = load_model(path_encoder)
    decoder = load_model(path_decoder)
    return encoder, decoder


# ============================================= 主流程相关 =======================================
# 比较 array1 相邻 array2 不相同的位数的矩阵
def compare_arrays(array1, array2):
    # 首先确保两个数组具有相同的形状
    if array1.shape != array2.shape:
        raise ValueError("两个数组的形状必须相同")
    # 结果是一个布尔数组，不同为 True，相同为 False
    comparison = np.not_equal(array1, array2)
    # 将布尔数组转换为整数数组，True 变为 1，False 变为 0
    result_matrix = comparison.astype(int)
    return result_matrix


# ================================================= 测试代码 ======================================================
def create_beams_phaseRads_phaseBits(bit_num):
    """ 生成波束指向和对应码阵 """
    beams = [(20, 0), (20, 60), (20, 120), (20, 180)]
    phaseRads = []
    phaseBits = []
    for beam in beams:
        # 生成码阵
        phaseRad = point_2_phase(theta0=beam[0], phi0=beam[1])
        phaseRad = np.rad2deg(phaseRad)
        # 码阵bit化
        phaseBit, phaseDeg = phase_2_bit(phaseRad, bit_num)
        # 记录
        phaseRads.append(phaseRad)
        phaseBits.append(phaseBit)
    return beams, phaseRads, phaseBits

def test_multi_beam_simple_add():
    bit_num = 1
    encoder, decoder = load_models(path_encoder="../files/dissertation/chapter_3/tf-ae-bit-cnn2d-cbam/encoder.h5",
                                   path_decoder="../files/dissertation/chapter_3/tf-ae-bit-cnn2d-cbam/decoder.h5")

    # 生成多个波束指向和对应码阵
    beams, phaseRads, phaseBits = create_beams_phaseRads_phaseBits(bit_num)

    # PS-complex合成
    phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_complex_beam_n(phases=phaseRads, bit_num=bit_num)

    # 将phaseBits转换为适合卷积自编码器的格式
    enc_list = []
    for phaseBit in phaseBits:
        enc = encoder.predict(phaseBit.reshape(1, 64, 64, 1))
        enc_list.append(enc.reshape(16, 16))

    # fixme: 编码器直接叠加
    mix = np.sum(enc_list, axis=0)
    # mix = np.max(enc_list, axis=0)

    # 将mix转换为适合卷积自编码器的格式
    mix_dec = decoder.predict(mix.reshape(1, 16, 16, 1))

    # 将概率转换为0/1二进制值
    mix_dec_bit = (mix_dec > 0.5).astype(np.int32)

    # 将解码结果还原为原始形状
    phaseBit_dec = mix_dec_bit.reshape(64, 64)
    mix_arr = mix.reshape(16, 16)

    pattern_real, point_real, psll_real = phaseBit_2_pattern_point(phaseBit_mix, bit_num)
    pattern_dec, point_dec, psll_dec = phaseBit_2_pattern_point(phaseBit_dec, bit_num)
    diff = compare_arrays(phaseBit_mix, phaseBit_dec)

    # 比较编解码器叠加和PS-complex合成
    plot_images_3x2(phaseBit_mix, mix_arr, phaseBit_dec, pattern_real, diff, pattern_dec, "sum")
    # save_images_2x2("../files/dissertation/chapter_4/sum.jpg",
    #                 phaseBit_mix, "mix(PS-complex)", phaseBit_dec, "mix(sum)",
    #                 pattern_real, "pattern(PS-complex)", pattern_dec, "pattern(sum)", "method: encoded sum")
    # save_images_2x2("../files/dissertation/chapter_4/max.jpg",
    #                 phaseBit_mix, "mix(PS-complex)", phaseBit_dec, "mix(max)",
    #                 pattern_real, "pattern(PS-complex)", pattern_dec, "pattern(max)", "method: encoded max")


def test_2beam_ae():
    bit_num = 1
    encoder, decoder = load_models(path_encoder="../files/dissertation/chapter_3/tf-ae-bit-cnn2d-cbam/encoder.h5",
                                   path_decoder="../files/dissertation/chapter_3/tf-ae-bit-cnn2d-cbam/decoder.h5")

    # 生成多个波束指向和对应码阵
    beams, phaseRads, phaseBits = create_beams_phaseRads_phaseBits(bit_num)

    # PS-complex合成
    phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_complex_beam_n(phases=phaseRads, bit_num=bit_num)

    # 转换为适合卷积自编码器的格式
    X = phaseBit_mix.reshape(1, 64, 64, 1)  # 保持二维结构，并增加通道维度
    X = X.astype('float32')  # 保持0/1值

    # 编码合成后码阵
    X_enc = encoder.predict(X)

    # 解码合成后码阵
    X_dec = decoder.predict(X_enc)

    # 将概率转换为0/1二进制值
    X_dec_bit = (X_dec > 0.5).astype(np.int32)

    # 还原为原始形状
    X_dec_bit_arr = X_dec_bit.reshape(64, 64)
    X_ori_arr = X.reshape(64, 64)
    X_enc_arr = X_enc.reshape(16, 16)

    # 计算方向图
    pattern_orig, point_orig, psll_orig = phaseBit_2_pattern_point(X_ori_arr, bit_num)
    pattern_dec, point_dec, psll_dec = phaseBit_2_pattern_point(X_dec_bit_arr, bit_num)

    # 计算orig和dec区别
    diff = compare_arrays(X_ori_arr, X_dec_bit_arr)

    # 对比合成码阵编解码
    plot_images_3x2(X_ori_arr, X_enc_arr, X_dec_bit_arr, pattern_orig, diff, pattern_dec, f"accuracy:")



if __name__ == "__main__":
    # 测试双波束合成: 简单add和max
    test_multi_beam_simple_add()
    # 测试双波束编码
    # test_2beam_ae()